#!/usr/bin python3
# -*- coding:UTF-8 -*-
# Author: nigo
import pandas as pd


class Node:
    def __init__(self, name):
        self.name = name # 节点名称
        self.father = None # 父节点
        self.children = [] # 子节点

    def add_child(self, obj):
        self.children.append(obj) # 添加子节点

    def add_father(self, obj):
        self.father = obj # 添加父节点

    def terminal(self):
        """末级节点"""
        terminal_list = [] # 创建末级节点列表
        for node in self.children: # 遍历子节节点
            if node.children: # 若子节点有子节点
                terminal_list += node.terminal() # 迭代函数
            else:
                terminal_list.append(node.name) # 添加到末级节点列表
        return terminal_list


class Tree:
    def __init__(self, df):
        self.init(df) # 初始化

    def init(self, df):
        father_list = df.iloc[:, 0] # df第一列为父节点名称
        child_list = df.iloc[:, 1] # df第二列为子节点名称
        # 父节点中未出现在子节点的为顶级节点
        top = [i for i in father_list if i not in child_list]
        top = top[0]

        relation = {} # 层级对应关系
        for index, row in df.iterrows(): # 循环dataframe每一行
            father = row[0] # 父节点名称
            child = row[1] # 子节点名称
            if father in relation.keys():
                children = relation[father]
                children.append(child)
                relation[father] = children
            else:
                relation[father] = [child]
        self.relation = relation
        self.top = self.make_tree(top)

    def make_tree(self, top):
        """根据顶级节点建立树状关系"""
        father_node = Node(top) # 创建父节点
        for child in self.relation[top]: # 遍历父节点对应的所有下一层级
            if child in self.relation.keys(): # 如果child有下一层级
                child_node = self.make_tree(child) # 调用迭代函数
                child_node.add_father(father_node) # 给子节点添加父节点
                father_node.add_child(child_node) # 给父节点添加子节点
            else: # 如果child没有下一层级
                child_node = Node(child) # 创建子节点
                child_node.add_father(father_node) # 给子节点添加父节点
                father_node.add_child(child_node) # 给父节点添加子节点
        return father_node

    def terminal(self):
        return self.top.terminal()


if __name__ == '__main__':
    data = [
        ['A', 'B'],
        ['A', 'C'],
        ['B', 'D'],
        ['B', 'E'],
        ['C', 'F'],
    ]
    df = pd.DataFrame(data)
    tree = Tree(df)
    print(tree.top.name)
    print(tree.top.terminal())
    print(tree.relation)
